import numpy as np

from osc03.static.constants import Constants


def pert_samples(a, b, c, size, rng):
    mu = (a + 4*b + c) / 6
    # 这个公式 不对
    alpha = ((mu - a) * (2*b - a - c)) / ((b - mu) * (c - a))
    beta = (alpha * (c - mu)) / (mu - a)
    return a + (c - a) * rng.beta(alpha, beta, size=size)

# 参数定义
a_std = Constants.STANDARD_FLOOR_ERECTION_LEFT
b_std = Constants.STANDARD_FLOOR_ERECTION_MODE
c_std = Constants.STANDARD_FLOOR_ERECTION_RIGHT

a_spec = a_std + 1
b_spec = b_std + 1
c_spec = c_std + 1

params = [(a_spec, b_spec, c_spec)] * 2 + [(a_std, b_std, c_std)] * 15 + [(a_spec, b_spec, c_spec)]

# 生成500组样本 (18×10,000矩阵)
rng = np.random.default_rng(seed=42)
samples = np.array([
    pert_samples(a, b, c, size=100, rng=rng)
    for a, b, c in params
])


def clip_by_row(samples, special_rows, a_spec, c_spec, a_std, c_std):
    """
    按行差异化裁剪矩阵

    参数:
        samples: 输入矩阵(17x10)
        special_rows: 特殊行索引列表(如[0,1,17])
        a_spec, c_spec: 特殊行边界
        a_std, c_std: 标准行边界
    """
    bounds = np.full((samples.shape[0], 2), [a_std, c_std])
    bounds[special_rows] = [a_spec, c_spec]
    return np.clip(samples, bounds[:, 0:1], bounds[:, 1:2])

samples = clip_by_row(samples, special_rows=[0,1,17], a_spec=a_spec,c_spec=c_spec,a_std=a_std,c_std=c_std)

import numpy as np
import matplotlib.pyplot as plt
import scipy.stats as stats

def check_pert_distribution(samples, a, b, c, row_idx=0, plot=True):
    """
    Validate if PERT-generated samples follow a Beta distribution with 4 diagnostic plots.

    Args:
        samples: Generated sample matrix (n_rows x n_samples)
        a, b, c: Optimistic, most likely, pessimistic times
        row_idx: Row index to validate
        plot: Whether to generate visual plots
    """
    data = samples[row_idx]
    mu = (a + 4 * b + c) / 6
    alpha = ((mu - a) * (2 * b - a - c)) / ((b - mu) * (c - a))
    beta_val = (alpha * (c - mu)) / (mu - a)

    # Theoretical distribution
    theoretical_dist = stats.beta(alpha, beta_val, loc=a, scale=c - a)
    x = np.linspace(a, c, 200)
    pdf = theoretical_dist.pdf(x)
    cdf = theoretical_dist.cdf(x)

    if plot:
        fig, axs = plt.subplots(2, 2, figsize=(15, 12))
        fig.suptitle(f'PERT Distribution Validation (Row {row_idx})\n'
                     f'a={a}, b={b}, c={c} | α={alpha:.2f}, β={beta_val:.2f}',
                     fontsize=14, y=1.02)

        # 1.1 Histogram vs PDF (Top-Left)
        axs[0, 0].hist(data, bins=30, density=True, alpha=0.6,
                       color='steelblue', label='Sampled Data')
        axs[0, 0].plot(x, pdf, 'r-', lw=2, label='Theoretical PDF')
        axs[0, 0].set_title('Density Comparison')
        axs[0, 0].set_xlabel('Time')
        axs[0, 0].set_ylabel('Probability Density')
        axs[0, 0].legend()
        axs[0, 0].grid(True, alpha=0.3)

        # 1.2 Q-Q Plot (Top-Right)
        stats.probplot(data, dist=theoretical_dist, plot=axs[0, 1])
        axs[0, 1].set_title('Quantile-Quantile Plot')
        axs[0, 1].get_lines()[0].set_markerfacecolor('steelblue')
        axs[0, 1].get_lines()[1].set_color('red')
        axs[0, 1].grid(True, alpha=0.3)

        # 1.3. CDF Comparison (Bottom-Left)
        ecdf = np.sort(data), np.arange(1, len(data) + 1) / len(data)
        axs[1, 0].plot(*ecdf, 'b-', lw=2, label='Empirical CDF')
        axs[1, 0].plot(x, cdf, 'r--', lw=2, label='Theoretical CDF')
        axs[1, 0].set_title('Cumulative Distribution Comparison')
        axs[1, 0].set_xlabel('Time')
        axs[1, 0].set_ylabel('Cumulative Probability')
        axs[1, 0].legend()
        axs[1, 0].grid(True, alpha=0.3)

        # 1.4 Boxplot with Theoretical Range (Bottom-Right)
        axs[1, 1].boxplot(data, vert=False, patch_artist=True,
                          boxprops=dict(facecolor='steelblue', alpha=0.6))
        axs[1, 1].axvline(a, color='green', linestyle='--', label='Optimistic (a)')
        axs[1, 1].axvline(b, color='orange', linestyle='--', label='Most Likely (b)')
        axs[1, 1].axvline(c, color='red', linestyle='--', label='Pessimistic (c)')
        axs[1, 1].set_title('Data Range vs PERT Estimates')
        axs[1, 1].set_xlabel('Time')
        axs[1, 1].legend()
        axs[1, 1].grid(True, alpha=0.3)

        plt.tight_layout()
        plt.show()

    # 2. Kolmogorov-Smirnov Test
    ks_stat, p_value = stats.kstest(data, theoretical_dist.cdf)
    print(f"Row {row_idx} K-S Test Results:")
    print(f"  Statistic: {ks_stat:.4f}, p-value: {p_value:.4f}")
    print(f"  Alpha: {alpha:.2f}, Beta: {beta_val:.2f}")
    if p_value > 0.05:
        print("  ✅ Fail to reject H0 (data likely follows Beta distribution)")
    else:
        print("  ❌ Reject H0 (data may NOT follow Beta distribution)")

    return {
        'alpha': alpha,
        'beta': beta_val,
        'ks_stat': ks_stat,
        'p_value': p_value
    }


# Validate special and standard rows
# print("===== Validating Special Rows =====")
# check_pert_distribution(samples, a_spec, b_spec, c_spec, row_idx=0)
#
# print("\n===== Validating Standard Rows =====")
# check_pert_distribution(samples, a_std, b_std, c_std, row_idx=2)


# Optional: Validate all rows
def validate_all_rows(samples, params):
    results = []
    for i, (a, b, c) in enumerate(params):
        print(f"\n=== Row {i} (a={a}, b={b}, c={c}) ===")
        res = check_pert_distribution(samples, a, b, c, row_idx=i, plot=True)
        results.append(res)
    return results

# Uncomment to run full validation (may produce verbose output)
params_all = [(a_spec, b_spec, c_spec)] * 2 + [(a_std, b_std, c_std)] * 15 + [(a_spec, b_spec, c_spec)]
validate_all_rows(samples, params_all)